# Author: Zylo117

"""
Simple Inference Script of EfficientDet-Pytorch
"""
import time
import torch
from tqdm import tqdm
import glob
from torch.backends import cudnn
from matplotlib import colors
import os
from backbone import EfficientDetBackbone
import cv2
import numpy as np
import os.path as osp
from efficientdet.utils import BBoxTransform, ClipBoxes
from utils.utils import preprocess, invert_affine, postprocess, STANDARD_COLORS, standard_to_bgr, get_index_label, plot_one_box

head_only = True
compound_coef = 0
result_dir = osp.join('results/head_only_D{}'.format(compound_coef)) if head_only else osp.join('results/D{}'.format(compound_coef))
if not osp.exists(result_dir):
    os.makedirs(result_dir)
force_input_size = None  # set None to use default size
test_dev_root = '/media/hp208/4t/zhaoxingjie/project_graduation/data/VisDrone2020/VisDrone2019-DET-test-dev/images/*'
img_paths = glob.glob(test_dev_root)

weight_path = 'logs/visdrone2020/efficientdet-d0_5_3000.pth'

# replace this part with your project's anchor config
anchor_ratios = [(1.0, 1.0), (1.4, 0.7), (0.7, 1.4)]
anchor_scales = [2 ** 0, 2 ** (1.0 / 3.0), 2 ** (2.0 / 3.0)]

threshold = 0.2
iou_threshold = 0.2

use_cuda = True
use_float16 = False
cudnn.fastest = True
cudnn.benchmark = True

obj_list = ['pedestrian',
           'person',
           'bicycle',
           'car', 'van', 'truck','tricycle','awning-tricycle','bus',
           'motor']


color_list = standard_to_bgr(STANDARD_COLORS)
# tf bilinear interpolation is different from any other's, just make do
input_sizes = [512, 640, 768, 896, 1024, 1280, 1280, 1536]
input_size = input_sizes[compound_coef] if force_input_size is None else force_input_size

model = EfficientDetBackbone(compound_coef=compound_coef, num_classes=len(obj_list),
                             ratios=anchor_ratios, scales=anchor_scales)
model.load_state_dict(torch.load(weight_path))
# model.requires_grad_(False)
model.eval()

if use_cuda:
    model = model.cuda()
if use_float16:
    model = model.half()



def display(preds, imgs, imshow=True, imwrite=False):
    for i in range(len(imgs)):
        if len(preds[i]['rois']) == 0:
            continue

        for j in range(len(preds[i]['rois'])):
            x1, y1, x2, y2 = preds[i]['rois'][j].astype(np.int)
            obj = obj_list[preds[i]['class_ids'][j]]
            score = float(preds[i]['scores'][j])
            plot_one_box(imgs[i], [x1, y1, x2, y2], label=obj,score=score,color=color_list[get_index_label(obj, obj_list)])


        if imshow:
            cv2.imshow('img', imgs[i])
            cv2.waitKey(0)

        if imwrite:
            cv2.imwrite('test/img_inferred_d{}_this_repo_{}.jpg'.format(compound_coef, i), imgs[i])


def write_text(res_path, preds):
    with open(res_path, 'w') as f:
        if len(preds[0]['rois']) == 0:
            return
        for j in range(len(preds[0]['rois'])):
            x1, y1, x2, y2 = preds[0]['rois'][j].astype(np.int)
            obj = preds[0]['class_ids'][j] + 1
            score = float(preds[0]['scores'][j])
            string = '{},{},{},{},{},{},0,0\n'.format(x1, y1, x2-x1, y2-y1, score, obj)
            # print(string)
            f.write(string)


def run():
    with torch.no_grad():
        for img_path in tqdm(img_paths):

            ori_imgs, framed_imgs, framed_metas = preprocess(img_path, max_size=input_size)
            x = torch.stack([torch.from_numpy(fi).cuda() for fi in framed_imgs], 0)

            x = x.to(torch.float32 if not use_float16 else torch.float16).permute(0, 3, 1, 2)
            features, regression, classification, anchors = model(x)

            regressBoxes = BBoxTransform()
            clipBoxes = ClipBoxes()

            out = postprocess(x,
                              anchors, regression, classification,
                              regressBoxes, clipBoxes,
                              threshold, iou_threshold)
            out = invert_affine(framed_metas, out)
            # display(out, ori_imgs, imshow=True, imwrite=False)
            res_path = osp.join(result_dir, osp.basename(img_path).replace('jpg', 'txt'))
            write_text(res_path, out)

if __name__=='__main__':
    run()

